import torch.nn as nn
from transformers import AutoModel
import torch.nn.functional as F

class MyBert(nn.Module):
    def __init__(self, opt):
        super(MyBert, self).__init__()
        self.dropout = nn.Dropout(opt.dropout)
        self.bert = AutoModel.from_pretrained(opt.plm, output_hidden_states = False) # 输出768维
        self.fc = nn.Sequential(nn.Linear(768, 283), nn.Dropout(opt.dropout), nn.Linear(283, 3))

    def forward(self, inputs):
        hidden_feats, cls_feats = self.bert(
            inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            token_type_ids=inputs['token_type_ids'],
            return_dict=False)
        out = F.relu(cls_feats)
        out = self.dropout(out)
        out = self.fc(out)

        return out